
// VS_BEGIN
#version 330 core
layout (location = 0) in vec3 aPos;
layout (location = 1) in vec2 aCoord;
layout (location = 2) in vec3 aNormal;
layout (location = 3) in vec3 aTangent;
layout (location = 4) in vec3 aBitangent;

uniform mat4 u_model;

out vec2 TexCoords;

layout (std140) uniform EngineUBO
{
    mat4 u_view;
    mat4 u_project;
    vec3 u_camPos;
    vec3 u_camInfo;
    mat4 u_inversePV;
    float u_time;
};

void main() {
    TexCoords = aCoord;
    gl_Position = vec4(aPos, 1.0f);
}
// VS_END
// FS_BEGIN
#version 330 core 

const float PI = 3.14159265359;
const float V = 2.0f / 3.0f * PI;

// Light Packet:
// position.x position.y position.z type(PointLight: 0; DirectLight: 1; AmbientSphereLight: 2)
// color.r    color.g    color.b    strength
// direct.x   direct.y   direct.z   unuse
// unuse      unuse      unuse      unuse
layout (std140) uniform LightUBO
{
    mat4 u_lights[64];
    int u_lightNum;
};

// u_camInfo Packet:
// type(Perspective: 0, Ortho: 1) zNear zFar
layout (std140) uniform EngineUBO
{
    mat4 u_view;
    mat4 u_project;
    vec3 u_camPos;
    vec3 u_camInfo;
    mat4 u_inversePV;
    float u_time;
};

in vec2 TexCoords;

out vec4 FragColor;

uniform sampler2D u_lastMap;
uniform int u_frameCount;

uniform samplerCube u_irradianceMap;
uniform samplerCube u_prefilterMap;
uniform sampler2D   u_brdfLUT;

const float MAX_DIS = 1000.0f;

const vec3 SkyColor = vec3(0.01f);
// const vec3 SkyColor = vec3(0.08, 0.08, 0.25);

struct Material
{
    vec3 albedo;
    float roughness;
    vec3 irr;
    float energy;
};

struct Ray 
{
    vec3 ori;
    vec3 dir;
};

struct Sphere 
{
    vec3 c;
    float r;
    Material material;
};

struct Plane
{
    float y;
    Material material;
};

#define OBJS_SPHERE_NUM 3
Sphere spheres[OBJS_SPHERE_NUM];
#define OBJS_PLANE_NUM 1
Plane planes[OBJS_PLANE_NUM];
void InitObject() {
    spheres[0] = Sphere(vec3(0.0f),              10.0f, Material(vec3(1.0f, 1.0f, 0.0f), 0.5f, vec3(0.0f),              0.0f));
    spheres[1] = Sphere(vec3(0.0f, 10.0f, -13.0f), 10.0f,  Material(vec3(1.0f),          0.5f, vec3(1.0f),  1.0f));
    spheres[2] = Sphere(vec3(1.5f, 10.0f, 0.0f), 0.5f,  Material(vec3(0.0f, 0.0f, 1.0f), 0.2f, vec3(0.0f),              0.0f));
    planes[0] = Plane(10.0f, Material(vec3(0.0f, 1.0f, 0.0f), 0.5f, vec3(0.0f), 0.0f));
}

float SdSphere(vec3 p, Sphere sphere) {
    return length(p - sphere.c) - sphere.r;
}

float SdPlane(vec3 p, Plane plane) {
    return abs(p.y - plane.y);
}

float SDF(vec3 p) 
{
    float dis = 10000.0f;
    for (int i = 0; i < OBJS_SPHERE_NUM; ++i) {
        float dis1 = SdSphere(p, spheres[i]);
        dis = min(dis, dis1);
    }
    for (int i = 0; i < OBJS_PLANE_NUM; ++i) {
        float dis1 = SdPlane(p, planes[i]);
        dis = min(dis, dis1);
    }
    return dis;
}

Material GetMaterial(vec3 p) {
    for (int i = 0; i < OBJS_SPHERE_NUM; ++i) {
        float dis = SdSphere(p, spheres[i]);
        if (dis < 0.0001f) {
            return spheres[i].material;
        }    
    }
    for (int i = 0; i < OBJS_PLANE_NUM; ++i) {
        float dis = SdPlane(p, planes[i]);
        if (dis < 0.0001f) {
            return planes[i].material;
        }   
    }
}

// 使用更好的随机数函数
uint seed = uint(u_time * 1000.0) + uint(gl_FragCoord.x) * 1973u + uint(gl_FragCoord.y) * 9277u;

uint wang_hash(inout uint seed) {
    seed = (seed ^ 61u) ^ (seed >> 16u);
    seed *= 9u;
    seed = seed ^ (seed >> 4u);
    seed *= 0x27d4eb2du;
    seed = seed ^ (seed >> 15u);
    return seed;
}

float rand(inout uint seed) {
    return float(wang_hash(seed)) / 4294967296.0;
}

vec3 cosineSampleHemisphere(vec3 normal, inout uint seed) {
    float r1 = rand(seed);
    float r2 = rand(seed);
    
    float phi = 2.0 * PI * r1;
    float cosTheta = sqrt(1.0 - r2);
    float sinTheta = sqrt(r2);
    
    vec3 hemisphere = vec3(cos(phi) * sinTheta, sin(phi) * sinTheta, cosTheta);
    
    // 转换到法线空间
    vec3 tangent = normalize(cross(normal, vec3(0, 1, 0)));
    if (length(tangent) < 0.001) 
        tangent = normalize(cross(normal, vec3(1, 0, 0)));
    vec3 bitangent = cross(normal, tangent);
    
    return tangent * hemisphere.x + bitangent * hemisphere.y + normal * hemisphere.z;
}

vec3 CalcNormal(vec3 p) {
    const float eps = 0.001; // 一个很小的数，用于差分
    const vec2 h = vec2(eps, 0);

    // 中心差分法计算偏导数
    float dx = SDF(p + h.xyy) - SDF(p - h.xyy);
    float dy = SDF(p + h.yxy) - SDF(p - h.yxy);
    float dz = SDF(p + h.yyx) - SDF(p - h.yyx);

    vec3 normal = vec3(dx, dy, dz) / (2.0 * eps);
    return normalize(normal);
}

bool IsHit(Ray ray, out vec3 p) {
    float totalDis = 0.0f;
    while (totalDis < MAX_DIS) {
        p = ray.ori + ray.dir * totalDis;
        float dis = SDF(p);
        if (dis < 0.0001f) {
            return true;
        }
        totalDis += dis;
    }
    return false;
}

vec3 RayMatch(Ray ray, int bounce) 
{
    vec3 light = vec3(0.0f);
    vec3 contribution = vec3(1.0f);
    for (int i = 0; i <= bounce; ++i) {
        uint raySeed = seed + uint(bounce) * 7919u + uint(i) * 104729u;
        vec3 p;
        if (IsHit(ray, p)) {
            Material m = GetMaterial(p);
            light += contribution * m.irr * m.energy;
            contribution *= m.albedo;
            vec3 normal = CalcNormal(p);
            ray.ori = p + normal * 0.001f;
            ray.dir = normalize(cosineSampleHemisphere(normal, raySeed));
        } else {
            light += SkyColor * contribution;
            break;
        }
    }
    return clamp(light, vec3(0.0f), vec3(1.0f));
}

void main() {
    InitObject();
    vec2 uv = TexCoords;
    vec3 NDC_pos = vec3(uv * 2.0f - 1.0f, -1.0f);
    vec4 world_pos_w = u_inversePV * vec4(NDC_pos, 1.0f);
    vec3 world_pos = world_pos_w.xyz / world_pos_w.w;
    Ray ray;
    ray.ori = u_camPos;
    ray.dir = normalize(world_pos - u_camPos);

    vec3 color = RayMatch(ray, 5);

    // 蒙特卡洛积分
    if (u_frameCount > 1) {
        FragColor = vec4((texture(u_lastMap, uv).rgb * float(u_frameCount - 1) + color * V) / float(u_frameCount), 1.0f);
    } else {
        FragColor = vec4(color * V / float(u_frameCount), 1.0f);
    }
}
// FS_END
